#include "onnx_parser.h"

namespace weilin {

OnnxParse::OnnxParse(uint32_t flag, std::size_t workspaceSize)
    : isParseSuccess(false) {
  mBuilder = std::unique_ptr<nvinfer1::IBuilder>(
      nvinfer1::createInferBuilder(nv_logger));
  mNetwork = std::unique_ptr<nvinfer1::INetworkDefinition>(
      mBuilder->createNetworkV2(flag));
  mParser = std::unique_ptr<nvonnxparser::IParser>(
      nvonnxparser::createParser(*mNetwork, nv_logger));
  mConfig = std::shared_ptr<nvinfer1::IBuilderConfig>(
      mBuilder->createBuilderConfig());
  mConfig->setMaxWorkspaceSize(workspaceSize);
}

void OnnxParse::parse(const std::string& file) {
  mParser->parseFromFile(file.c_str(),
                         (int)nvinfer1::ILogger::Severity::kWARNING);
  if (mParser->getNbErrors() != 0) {
    isParseSuccess = false;
    for (int i = 0; i < mParser->getNbErrors(); ++i) {
      spdlog::error(mParser->getError(i)->desc());
    }

    return;
  }
  mModel = std::shared_ptr<nvinfer1::IHostMemory>(
      mBuilder->buildSerializedNetwork(*mNetwork, *mConfig));
  nvinfer1::IErrorRecorder* errorRecorder = mBuilder->getErrorRecorder();
  if (errorRecorder != nullptr) {
    for (int i = 0; i < errorRecorder->getNbErrors(); ++i) {
      spdlog::error(errorRecorder->getErrorDesc(i));
    }
    isParseSuccess = false;
    return;
  }
  isParseSuccess = true;
}

void OnnxParse::save_model(const std::string& file) {
  std::ofstream of(file);
  of.write(static_cast<char*>(mModel->data()), mModel->size());
  of.close();
}

}  // namespace weilin
